#!/usr/bin/env python3
"""
Transformer视频生成器测试套件
验证视频生成模型的正确性和性能
"""
import torch
import pytest
import sys
import os
import numpy as np

# Add the models directory to the path
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', 'models'))

from models.video_generator import VideoGenerator

class TestVideoGenerator:
    def setup_method(self):
        """Setup method for each test"""
        self.generator = VideoGenerator(
            frame_size=(32, 32),
            patch_size=8,
            num_frames=4,
            embedding_dim=128,
            nhead=4,
            num_layers=2,
            num_channels=3,
            max_patches=256
        )
    
    def test_generator_initialization(self):
        """Test that the video generator initializes correctly"""
        assert self.generator is not None
        assert isinstance(self.generator, VideoGenerator)
        assert hasattr(self.generator, 'model')
        
        # Test attributes
        assert self.generator.frame_size == (32, 32)
        assert self.generator.patch_size == 8
        assert self.generator.num_frames == 4
        assert self.generator.embedding_dim == 128
        assert self.generator.num_channels == 3
    
    def test_patch_calculation(self):
        """Test patch calculation logic"""
        assert self.generator.patches_per_row == 4  # 32 // 8
        assert self.generator.patches_per_col == 4  # 32 // 8
        assert self.generator.patches_per_frame == 16  # 4 * 4
        assert self.generator.total_patches == 64  # 16 * 4
    
    def test_video_to_patches_conversion(self):
        """Test video to patches conversion"""
        # Create test video tensor
        batch_size = 2
        video_tensor = torch.randn(batch_size, self.generator.num_frames, 
                                  self.generator.num_channels, 
                                  self.generator.frame_size[0], 
                                  self.generator.frame_size[1])
        
        patches = self.generator._video_to_patches(video_tensor)
        
        # Check output shape
        expected_shape = (batch_size, self.generator.total_patches, 
                         self.generator.patch_dim)
        assert patches.shape == expected_shape
        assert not torch.isnan(patches).any()
        assert not torch.isinf(patches).any()
    
    def test_patches_to_video_conversion(self):
        """Test patches to video conversion"""
        # Create test patches tensor
        batch_size = 2
        patches = torch.randn(batch_size, self.generator.total_patches, 
                             self.generator.patch_dim)
        
        video = self.generator._patches_to_video(patches)
        
        # Check output shape
        expected_shape = (batch_size, self.generator.num_frames, 
                         self.generator.num_channels,
                         self.generator.frame_size[0], 
                         self.generator.frame_size[1])
        assert video.shape == expected_shape
        assert not torch.isnan(video).any()
        assert not torch.isinf(video).any()
    
    def test_round_trip_conversion(self):
        """Test that video -> patches -> video conversion is lossless"""
        # Create original video
        batch_size = 1
        original_video = torch.randn(batch_size, self.generator.num_frames, 
                                   self.generator.num_channels,
                                   self.generator.frame_size[0], 
                                   self.generator.frame_size[1])
        
        # Convert to patches and back
        patches = self.generator._video_to_patches(original_video)
        reconstructed_video = self.generator._patches_to_video(patches)
        
        # Should have same shape
        assert original_video.shape == reconstructed_video.shape
        
        # Values should be close (allowing for numerical precision)
        torch.testing.assert_close(original_video, reconstructed_video, rtol=1e-4, atol=1e-6)
    
    def test_generate_method(self):
        """Test the generate method with simple input"""
        # Create input video with correct shape: (batch_size, num_frames, channels, height, width)
        batch_size = 1
        input_video = torch.randn(batch_size, self.generator.num_frames, self.generator.num_channels,
                                 self.generator.frame_size[0], 
                                 self.generator.frame_size[1])
        
        # Test generation
        generated_video = self.generator.generate(input_video, num_steps=1)
        
        # Check output shape (should be same as input)
        expected_shape = (batch_size, self.generator.num_frames+1, self.generator.num_channels,
                         self.generator.frame_size[0], 
                         self.generator.frame_size[1])
        assert generated_video.shape == expected_shape
        
        # Check value range (should be normalized to [0, 1])
        assert generated_video.min() >= 0.0
        assert generated_video.max() <= 1.0
    
    def test_generate_continuation(self):
        """Test video continuation generation"""
        # Create input video
        batch_size = 1
        input_video = torch.randn(batch_size, self.generator.num_frames, self.generator.num_channels,
                                 self.generator.frame_size[0], 
                                 self.generator.frame_size[1])
        
        # Test continuation
        continued_video = self.generator.generate_continuation(input_video, continuation_frames=2)
        
        # Check output shape (input frames + continuation frames)
        expected_shape = (batch_size, self.generator.num_frames + 2, self.generator.num_channels,
                         self.generator.frame_size[0], 
                         self.generator.frame_size[1])
        assert continued_video.shape == expected_shape
    
    def test_edge_cases(self):
        """Test edge cases"""
        # Test with single frame
        self.generator = VideoGenerator(
            frame_size=(32, 32),
            patch_size=8,
            num_frames=1,
            embedding_dim=128,
            nhead=4,
            num_layers=2,
            num_channels=3,
            max_patches=256
        )
        single_frame = torch.randn(1, 1, self.generator.num_channels,
                                  self.generator.frame_size[0], 
                                  self.generator.frame_size[1])
        result = self.generator.generate(single_frame, num_steps=1)
        assert result.shape[0] == single_frame.shape[0]
        assert result.shape[1] == single_frame.shape[1] + 1
        assert result.shape[2] == single_frame.shape[2]
        assert result.shape[3] == single_frame.shape[3]
        assert result.shape[4] == single_frame.shape[4]
        
        # Test with zero frames (should handle gracefully)
        with pytest.raises(Exception):
            zero_frames = torch.randn(1, 0, self.generator.num_channels,
                                     self.generator.frame_size[0], 
                                     self.generator.frame_size[1])
            self.generator.generate(zero_frames, num_steps=1)
    
    def test_temperature_parameter(self):
        """Test temperature parameter in generation"""
        input_video = torch.randn(1, self.generator.num_frames, self.generator.num_channels,
                                self.generator.frame_size[0], 
                                self.generator.frame_size[1])
        
        # Test with different temperatures
        video_low_temp = self.generator.generate(input_video, num_steps=1, temperature=0.1)
        video_high_temp = self.generator.generate(input_video, num_steps=1, temperature=2.0)
        
        # Both should produce valid output
        assert video_low_temp.shape[0] == input_video.shape[0]
        assert video_low_temp.shape[1] == input_video.shape[1] + 1
        assert video_low_temp.shape[2] == input_video.shape[2]
        assert video_low_temp.shape[3] == input_video.shape[3]
        assert video_low_temp.shape[4] == input_video.shape[4]
        assert video_high_temp.shape[0] == input_video.shape[0]
        assert video_high_temp.shape[1] == input_video.shape[1] + 1
        assert video_high_temp.shape[2] == input_video.shape[2]
        assert video_high_temp.shape[3] == input_video.shape[3]
        assert video_high_temp.shape[4] == input_video.shape[4]

if __name__ == "__main__":
    pytest.main([__file__])